#!/usr/bin/env python
# -*- coding: UTF-8 -*-
'''
@Project ：V2 
@File    ：t4.py
@IDE     ：PyCharm 
@Author  ：郭星
@Date    ：2025/9/15 21:41

'''
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
多种算法求解车间调度问题
"""
"""


说明与二次开发指南

模块化设计

两类问题各自封装：ProblemFJSP 与 ProblemHFS，都提供 decode(keys)->(makespan, assign) 与 plot_gantt(assign, title)

算法封装为函数：run_ga/ run_abc/ run_foa/ run_gwo，输入 WrapperObjective（包含 decoder 与 dim），返回 (best_x, best_f, history, assign)

统一编码 Random Keys

FJSP：随机键个数等于总工序数；按键值升序逐步调度，每个工序在其可行机上选择最早完工

HFS：随机键个数等于工件数；按键值升序得到工件顺序；每阶段在并行机上选择最早可开工

这种统一编码便于在同一套算法框架里复用

收敛曲线：各算法的 history（每代/次迭代最优完工期）用 plot_convergence 绘制

甘特图：ProblemFJSP.plot_gantt 与 ProblemHFS.plot_gantt 分别绘制

扩展点

想要多目标（如同时考虑在制品/延迟），可以把 decoder 返回多个指标，改写 WrapperObjective.fitness 为加权或分层排序

想要重启/并行：把 run_* 改成支持多个独立种群或多次独立运行取最好

想要固定机器分配作为基因（而非贪心选机），可以把 keys 拆为两段：优先级段 + 机台码段；在 decode 中读取机台码（取模到可行集合）代替贪心选机

如果你希望我把实例规模直接改成“FJSP 20×4×5 / HFS 若干阶段×并行机”并保存结果到文件（CSV/PNG），或者把这四种算法做成一个统一的类方便在外部脚本里 import 调用，我也可以一并给你整理好。
Metaheuristics for FJSP & HFS (GA / ABC / FOA / GWO)
- Unified random-keys encoding
- Convergence curves + Gantt charts
- Modular, reusable API

Usage:
  pip install numpy matplotlib
  python meta_fjsp_hfs.py
"""

import math
import random
from dataclasses import dataclass
from typing import Dict, List, Tuple, Callable, Optional
import numpy as np
import matplotlib.pyplot as plt

# ------------------------------
# Utilities
# ------------------------------

def set_seed(seed: int = 42):
    np.random.seed(seed)
    random.seed(seed)

def clamp01(x: np.ndarray) -> np.ndarray:
    return np.minimum(1.0, np.maximum(0.0, x))

def rand_keys(dim: int) -> np.ndarray:
    return np.random.rand(dim).astype(np.float32)

def tournament_select_idx(f: np.ndarray, k: int = 3) -> int:
    idx = np.random.randint(0, len(f), size=k)
    best = idx[np.argmin(f[idx])]
    return int(best)

# ------------------------------
# FJSP definition & decoder
# ------------------------------

@dataclass
class Operation:
    machine_time: Dict[int, int]  # feasible machine -> time

@dataclass
class JobFJSP:
    ops: List[Operation]
    due: float = 0.0  # optional

@dataclass
class ProblemFJSP:
    jobs: List[JobFJSP]
    m: int  # number of machines

    @property
    def n_jobs(self): return len(self.jobs)

    @property
    def n_ops_total(self): return sum(len(j.ops) for j in self.jobs)

    def decode(self, keys: np.ndarray):
        """
        Random-keys decoding for FJSP:
        - Each (job, op) has a key; schedule in ascending key order but obey precedence.
        - For each ready operation, choose machine that yields earliest completion.
        Return: makespan, assign dict: (j,k)->(m,start,end)
        """
        n_jobs, M = self.n_jobs, self.m
        assert len(keys) == self.n_ops_total
        # mapping from linear index -> (j, k)
        mapping = []
        for j, job in enumerate(self.jobs):
            for k in range(len(job.ops)):
                mapping.append((j, k))
        # keys per (j,k)
        key_per_jk = { (j,k): keys[i] for i,(j,k) in enumerate(mapping) }

        next_op = np.zeros(n_jobs, dtype=np.int32)
        job_ready = np.zeros(n_jobs, dtype=np.float32)
        mach_ready = np.zeros(M, dtype=np.float32)
        makespan = 0.0
        done_ops = 0
        total_ops = self.n_ops_total
        assign: Dict[Tuple[int,int], Tuple[int,float,float]] = {}

        # To efficiently pick among ready ops by key:
        while done_ops < total_ops:
            # collect ready ops (next_op[j] not finished)
            ready = []
            for j in range(n_jobs):
                k = int(next_op[j])
                if k < len(self.jobs[j].ops):
                    ready.append((j, k, key_per_jk[(j,k)]))
            if not ready:
                # Should not happen, but jump time if stuck
                tmin = min(float(np.min(job_ready)), float(np.min(mach_ready)))
                job_ready[:] = np.maximum(job_ready, tmin)
                mach_ready[:] = np.maximum(mach_ready, tmin)
                continue

            # choose the smallest key among ready
            ready.sort(key=lambda x: x[2])
            j, k, _ = ready[0]
            op = self.jobs[j].ops[k]
            # choose machine with earliest completion
            best = None
            best_end = float("inf")
            best_start = 0.0
            best_m = None
            for mm, pt in op.machine_time.items():
                start = max(float(job_ready[j]), float(mach_ready[mm]))
                end = start + float(pt)
                if (end < best_end) or (abs(end - best_end) < 1e-9 and start < best_start):
                    best_end = end
                    best_start = start
                    best = (mm, start, end)
                    best_m = mm
            # schedule
            mm, s, e = best
            job_ready[j] = e
            mach_ready[best_m] = e
            assign[(j, k)] = (best_m, s, e)
            makespan = max(makespan, e)
            next_op[j] += 1
            done_ops += 1

        return float(makespan), assign

    @staticmethod
    def random_instance(n_jobs=10, m=5, ops_per_job=4, seed=7):
        rng = random.Random(seed)
        jobs: List[JobFJSP] = []
        for _ in range(n_jobs):
            ops = []
            for _o in range(ops_per_job):
                k = 2 if m >= 2 else 1
                machines = sorted(rng.sample(range(m), k=k))
                mt = {mm: rng.randint(1, 9) for mm in machines}
                ops.append(Operation(mt))
            jobs.append(JobFJSP(ops=ops))
        return ProblemFJSP(jobs=jobs, m=m)

    # Gantt for FJSP
    def plot_gantt(self, assign: Dict[Tuple[int,int], Tuple[int,float,float]], title: str = "FJSP schedule"):
        bars = {mm: [] for mm in range(self.m)}
        for (j,k), (mm,s,e) in assign.items():
            bars[mm].append((s, e-s, f"J{j}-O{k}"))
        fig, ax = plt.subplots(figsize=(12,5))
        yticks, ylabels = [], []
        for idx, mm in enumerate(sorted(bars.keys())):
            items = sorted(bars[mm], key=lambda x: x[0])
            for (start, dur, label) in items:
                ax.barh(idx, dur, left=start)
                ax.text(start+dur/2, idx, label, ha="center", va="center", fontsize=8)
            yticks.append(idx); ylabels.append(f"M{mm}")
        ax.set_yticks(yticks); ax.set_yticklabels(ylabels)
        ax.set_xlabel("Time"); ax.set_title(title)
        plt.tight_layout(); plt.show()

# ------------------------------
# HFS definition & decoder
# ------------------------------

@dataclass
class Stage:
    machines: int
    proc_time: Dict[int, float]  # job_id -> processing time at this stage

@dataclass
class ProblemHFS:
    stages: List[Stage]  # length = #stages
    n_jobs: int

    @property
    def s(self): return len(self.stages)

    def decode(self, keys: np.ndarray):
        """
        Random-keys decoding for HFS:
        - keys define a permutation of jobs
        - For each stage, schedule jobs in that order assigning earliest available machine
        Return: makespan, assign dict: (job,stage)->(mach,start,end)
        """
        assert len(keys) == self.n_jobs
        order = np.argsort(keys).tolist()

        # For each stage keep machine ready times
        mach_ready = [np.zeros(stg.machines, dtype=np.float32) for stg in self.stages]
        job_ready = np.zeros(self.n_jobs, dtype=np.float32)
        assign: Dict[Tuple[int,int], Tuple[int,float,float]] = {}
        makespan = 0.0

        for stg_idx, stg in enumerate(self.stages):
            for job in order:
                pt = float(stg.proc_time[job])
                # pick machine with earliest completion
                best = None
                best_end = float("inf")
                m_idx_best = None
                for m_idx in range(stg.machines):
                    start = max(float(mach_ready[stg_idx][m_idx]), float(job_ready[job]))
                    end = start + pt
                    if end < best_end:
                        best_end = end
                        best = (m_idx, start, end)
                        m_idx_best = m_idx
                m_idx, s, e = best
                mach_ready[stg_idx][m_idx_best] = e
                job_ready[job] = e
                assign[(job, stg_idx)] = (m_idx, s, e)
                makespan = max(makespan, e)

        return float(makespan), assign

    @staticmethod
    def random_instance(n_jobs=12, stages=4, machines_per_stage=(2,2,2,2), seed=7):
        rng = random.Random(seed)
        stgs: List[Stage] = []
        for s_idx in range(stages):
            m = machines_per_stage[s_idx] if isinstance(machines_per_stage, (list,tuple)) else machines_per_stage
            pt = {j: rng.randint(1, 9) for j in range(n_jobs)}
            stgs.append(Stage(machines=m, proc_time=pt))
        return ProblemHFS(stages=stgs, n_jobs=n_jobs)

    def plot_gantt(self, assign: Dict[Tuple[int,int], Tuple[int,float,float]], title: str = "HFS schedule"):
        # Flatten machines across stages: (stage, machine)
        # Map each (stage, machine) to a y row
        rows = []
        for stg_idx, stg in enumerate(self.stages):
            for m_idx in range(stg.machines):
                rows.append((stg_idx, m_idx))
        row_index = {rows[i]: i for i in range(len(rows))}

        fig, ax = plt.subplots(figsize=(12, 0.8*len(rows)+3))
        yticks, ylabels = [], []
        # Collect bars per row
        per_row: Dict[int, List[Tuple[float,float,str]]] = {i: [] for i in range(len(rows))}
        for (job, stg_idx), (m_idx, s, e) in assign.items():
            r = row_index[(stg_idx, m_idx)]
            per_row[r].append((s, e-s, f"J{job}-S{stg_idx}"))
        for r in range(len(rows)):
            items = sorted(per_row[r], key=lambda x: x[0])
            for (start, dur, label) in items:
                ax.barh(r, dur, left=start)
                ax.text(start+dur/2, r, label, ha="center", va="center", fontsize=8)
            stg_idx, m_idx = rows[r]
            yticks.append(r); ylabels.append(f"S{stg_idx}-M{m_idx}")
        ax.set_yticks(yticks); ax.set_yticklabels(ylabels)
        ax.set_xlabel("Time"); ax.set_title(title)
        plt.tight_layout(); plt.show()

# ------------------------------
# Objective wrapper
# ------------------------------

class WrapperObjective:
    """
    Wraps a problem's decode(keys) -> (makespan, assign)
    Provides f(x) = makespan for algorithms; lower is better.
    """
    def __init__(self, decoder: Callable[[np.ndarray], Tuple[float, dict]], dim: int):
        self.decoder = decoder
        self.dim = dim

    def fitness(self, x: np.ndarray) -> float:
        x = clamp01(x)
        mk, _ = self.decoder(x)
        return mk

    def decode(self, x: np.ndarray):
        x = clamp01(x)
        return self.decoder(x)

# ------------------------------
# Metaheuristics
# ------------------------------

def run_ga(obj: WrapperObjective, iters=300, pop=50, elite=2, pc=0.9, pm=0.1, k=3):
    """
    Real-coded GA (random keys):
    - Tournament selection
    - BLX-α crossover (α=0.3) + uniform arithmetic fallback
    - Gaussian mutation (σ=0.1)
    """
    dim = obj.dim
    alpha = 0.3
    sigma = 0.1

    X = np.random.rand(pop, dim).astype(np.float32)
    F = np.array([obj.fitness(x) for x in X], dtype=np.float32)
    hist = [float(np.min(F))]
    best_idx = int(np.argmin(F))
    best_x = X[best_idx].copy()
    best_f = float(F[best_idx])

    for t in range(1, iters+1):
        # Elites
        idx_sorted = np.argsort(F)
        elites = X[idx_sorted[:elite]].copy()

        # Mating
        Y = []
        while len(Y) < pop - elite:
            p1 = X[tournament_select_idx(F, k)]
            p2 = X[tournament_select_idx(F, k)]
            c1, c2 = p1.copy(), p2.copy()
            if np.random.rand() < pc:
                lo = np.minimum(p1, p2)
                hi = np.maximum(p1, p2)
                range_ = hi - lo
                c_lo = lo - alpha * range_
                c_hi = hi + alpha * range_
                c1 = np.random.uniform(c_lo, c_hi).astype(np.float32)
                c2 = np.random.uniform(c_lo, c_hi).astype(np.float32)
            # mutation
            if np.random.rand() < pm:
                c1 += np.random.randn(dim).astype(np.float32) * sigma
            if np.random.rand() < pm:
                c2 += np.random.randn(dim).astype(np.float32) * sigma
            c1 = clamp01(c1); c2 = clamp01(c2)
            Y.append(c1)
            if len(Y) < pop - elite:
                Y.append(c2)

        X = np.vstack([elites] + Y)
        F = np.array([obj.fitness(x) for x in X], dtype=np.float32)

        cur_min = float(np.min(F))
        if cur_min < best_f:
            best_idx = int(np.argmin(F))
            best_x = X[best_idx].copy()
            best_f = cur_min
        hist.append(best_f)

    mk, assign = obj.decode(best_x)
    return best_x, best_f, hist, assign

def run_abc(obj: WrapperObjective, iters=300, pop=40, limit=20):
    """
    Artificial Bee Colony (continuous)
    - Employed: neighbor by phi*(xi - xk) on one dim
    - Onlooker: roulette by normalized (1/F)
    - Scout: reinit if trial >= limit
    """
    dim = obj.dim
    X = np.random.rand(pop, dim).astype(np.float32)
    F = np.array([obj.fitness(x) for x in X], dtype=np.float32)
    trial = np.zeros(pop, dtype=np.int32)
    hist = [float(np.min(F))]
    best_idx = int(np.argmin(F))
    best_x = X[best_idx].copy()
    best_f = float(F[best_idx])

    def neighbor(xi, xk, j):
        phi = np.random.uniform(-1, 1)
        v = xi.copy()
        v[j] = xi[j] + phi * (xi[j] - xk[j])
        return clamp01(v)

    for t in range(1, iters+1):
        # Employed
        for i in range(pop):
            k = np.random.choice([kk for kk in range(pop) if kk != i])
            j = np.random.randint(0, dim)
            v = neighbor(X[i], X[k], j)
            fv = obj.fitness(v)
            if fv < F[i]:
                X[i], F[i] = v, fv
                trial[i] = 0
            else:
                trial[i] += 1

        # Onlooker
        fit_inv = 1.0 / (F + 1e-9)
        prob = fit_inv / np.sum(fit_inv)
        for _ in range(pop):
            i = np.random.choice(np.arange(pop), p=prob)
            k = np.random.choice([kk for kk in range(pop) if kk != i])
            j = np.random.randint(0, dim)
            v = neighbor(X[i], X[k], j)
            fv = obj.fitness(v)
            if fv < F[i]:
                X[i], F[i] = v, fv
                trial[i] = 0
            else:
                trial[i] += 1

        # Scout
        for i in range(pop):
            if trial[i] >= limit:
                X[i] = np.random.rand(dim).astype(np.float32)
                F[i] = obj.fitness(X[i])
                trial[i] = 0

        # record
        cur_min = float(np.min(F))
        if cur_min < best_f:
            best_idx = int(np.argmin(F))
            best_x = X[best_idx].copy()
            best_f = cur_min
        hist.append(best_f)

    mk, assign = obj.decode(best_x)
    return best_x, best_f, hist, assign

def run_foa(obj: WrapperObjective, iters=300, swarm=30, step_init=0.2, step_decay=0.99):
    """
    Fruit Fly Optimization (vector form):
    - Maintain center; sample swarm candidates around center with step (Gaussian)
    - Move center to best; step decays each iteration
    """
    dim = obj.dim
    center = np.random.rand(dim).astype(np.float32)
    best_x = center.copy()
    best_f = obj.fitness(best_x)
    step = step_init
    hist = [best_f]

    for t in range(1, iters+1):
        cand = []
        for _ in range(swarm):
            v = center + np.random.randn(dim).astype(np.float32) * step
            v = clamp01(v)
            cand.append(v)
        cand = np.stack(cand, axis=0)
        F = np.array([obj.fitness(v) for v in cand], dtype=np.float32)
        idx = int(np.argmin(F))
        center = cand[idx].copy()
        if F[idx] < best_f:
            best_f = float(F[idx])
            best_x = center.copy()
        hist.append(best_f)
        step *= step_decay

    mk, assign = obj.decode(best_x)
    return best_x, best_f, hist, assign

def run_gwo(obj: WrapperObjective, iters=300, pop=40):
    """
    Grey Wolf Optimizer (continuous)
    """
    dim = obj.dim
    X = np.random.rand(pop, dim).astype(np.float32)
    F = np.array([obj.fitness(x) for x in X], dtype=np.float32)
    # alpha, beta, delta
    idx_sort = np.argsort(F)
    alpha, beta, delta = X[idx_sort[0]].copy(), X[idx_sort[1]].copy(), X[idx_sort[2]].copy()
    f_alpha, f_beta, f_delta = float(F[idx_sort[0]]), float(F[idx_sort[1]]), float(F[idx_sort[2]])
    hist = [f_alpha]

    for t in range(1, iters+1):
        a = 2.0 - 2.0 * (t / iters)  # linearly decreasing 2->0
        for i in range(pop):
            r1, r2 = np.random.rand(dim), np.random.rand(dim)
            A1 = 2*a*r1 - a
            C1 = 2*r2
            D_alpha = np.abs(C1*alpha - X[i])
            X1 = alpha - A1*D_alpha

            r1, r2 = np.random.rand(dim), np.random.rand(dim)
            A2 = 2*a*r1 - a
            C2 = 2*r2
            D_beta = np.abs(C2*beta - X[i])
            X2 = beta - A2*D_beta

            r1, r2 = np.random.rand(dim), np.random.rand(dim)
            A3 = 2*a*r1 - a
            C3 = 2*r2
            D_delta = np.abs(C3*delta - X[i])
            X3 = delta - A3*D_delta

            X[i] = clamp01((X1 + X2 + X3) / 3.0)

        F = np.array([obj.fitness(x) for x in X], dtype=np.float32)
        idx_sort = np.argsort(F)
        if F[idx_sort[0]] < f_alpha:
            alpha = X[idx_sort[0]].copy(); f_alpha = float(F[idx_sort[0]])
        if F[idx_sort[1]] < f_beta:
            beta = X[idx_sort[1]].copy(); f_beta = float(F[idx_sort[1]])
        if F[idx_sort[2]] < f_delta:
            delta = X[idx_sort[2]].copy(); f_delta = float(F[idx_sort[2]])
        hist.append(f_alpha)

    mk, assign = obj.decode(alpha)
    return alpha, f_alpha, hist, assign

# ------------------------------
# Runner & Plotting
# ------------------------------

def plot_convergence(hist: List[float], title: str):
    xs = np.arange(0, len(hist))
    plt.figure(figsize=(7,4))
    plt.plot(xs, hist)
    plt.xlabel("Iteration"); plt.ylabel("Best makespan")
    plt.title(title)
    plt.grid(True, alpha=0.3)
    plt.tight_layout(); plt.show()

def run_all_algorithms(problem_name: str, obj: WrapperObjective,
                       gantt_plotter: Callable[[dict, str], None],
                       label_prefix: str = "", iters=300):
    res = {}
    print(f"\n=== {problem_name}: GA ===")
    x, f, h, assign = run_ga(obj, iters=iters, pop=50)
    res["GA"] = (f, h, assign)
    plot_convergence(h, f"{label_prefix}GA convergence (best makespan)")
    gantt_plotter(assign, f"{label_prefix}GA best schedule (mk={f:.1f})")

    print(f"\n=== {problem_name}: ABC ===")
    x, f, h, assign = run_abc(obj, iters=iters, pop=40, limit=25)
    res["ABC"] = (f, h, assign)
    plot_convergence(h, f"{label_prefix}ABC convergence (best makespan)")
    gantt_plotter(assign, f"{label_prefix}ABC best schedule (mk={f:.1f})")

    print(f"\n=== {problem_name}: FOA ===")
    x, f, h, assign = run_foa(obj, iters=iters, swarm=50, step_init=0.25, step_decay=0.985)
    res["FOA"] = (f, h, assign)
    plot_convergence(h, f"{label_prefix}FOA convergence (best makespan)")
    gantt_plotter(assign, f"{label_prefix}FOA best schedule (mk={f:.1f})")

    print(f"\n=== {problem_name}: GWO ===")
    x, f, h, assign = run_gwo(obj, iters=iters, pop=40)
    res["GWO"] = (f, h, assign)
    plot_convergence(h, f"{label_prefix}GWO convergence (best makespan)")
    gantt_plotter(assign, f"{label_prefix}GWO best schedule (mk={f:.1f})")

    return res

# ------------------------------
# Demo main
# ------------------------------

def demo_fjsp():
    print("\n############################")
    print("# Demo: Flexible Job Shop  ")
    print("############################")
    # 小规模示例，运行较快；你可以改大如 n_jobs=20, m=5, ops_per_job=4
    fjsp = ProblemFJSP.random_instance(n_jobs=10, m=5, ops_per_job=4, seed=7)
    dim = fjsp.n_ops_total
    obj = WrapperObjective(fjsp.decode, dim=dim)
    res = run_all_algorithms("FJSP", obj, fjsp.plot_gantt, label_prefix="FJSP - ", iters=200)
    return fjsp, res

def demo_hfs():
    print("\n############################")
    print("# Demo: Hybrid Flow Shop   ")
    print("############################")
    # 混合流水车间：12 工件，4 阶段，每阶段 2 台并行机
    hfs = ProblemHFS.random_instance(n_jobs=12, stages=4, machines_per_stage=[2,2,2,2], seed=7)
    dim = hfs.n_jobs
    obj = WrapperObjective(hfs.decode, dim=dim)
    res = run_all_algorithms("HFS", obj, hfs.plot_gantt, label_prefix="HFS - ", iters=200)
    return hfs, res

def main():
    set_seed(42)
    fjsp, fjsp_res = demo_fjsp()
    hfs,  hfs_res  = demo_hfs()

if __name__ == "__main__":
    main()
